import os,pickle, gensim,time
import numpy as np
from collections import defaultdict

rootDir = os.path.dirname(os.getcwd())
isCodingAndDebug = __name__ == "__main__"

def get_words_matrix():
    # 获取词向量，词表
    matrix_str = np.loadtxt(os.path.join(rootDir, 'GloVe', 'vectors.txt'), dtype=str)
    words,matrix = matrix_str[:,0], matrix_str[:,1:].astype(float)
    word2id = dict(zip(words,range(len(words))))
    return words,word2id,matrix

dictSentence2KeyWordsStr = pickle.load(open(os.path.join(rootDir,"dictSentence2KeyWordsStr.pkl"),'rb'))
keyWordsStrList = dictSentence2KeyWordsStr.values()
dictID2Sentence = dict(zip(range(len(dictSentence2KeyWordsStr)),dictSentence2KeyWordsStr.keys()))

def get_keyWordsSentence_matrix(word2id,matrix,keyWordsStrList=keyWordsStrList):
    #获取一组解析词汇对应的平均向量。
    #输出维度：句子的个数 × 句向量个数
    len_words = matrix.shape[0]
    keyWordsStrMatrix = np.zeros((len(keyWordsStrList),len_words),int)
    normalizer = np.zeros((len(keyWordsStrList),1),int)

    func_listWords2listIDs = np.vectorize(lambda word: word2id[word] if word in word2id else word2id["<unk>"])

    def line2count(bundle):
        if isCodingAndDebug : print(bundle)
        ix, line = int(bundle[0]),bundle[1]
        ids = func_listWords2listIDs(line.split())

        def countVecUpdate(idxInWordsList):
            keyWordsStrMatrix[ix][idxInWordsList] += 1
        np.vectorize(countVecUpdate )(ids)

        # for idxInWordsList in ids:
        #     keyWordsStrMatrix[ix][idxInWordsList] += 1
        normalizer[ix]= len(ids)
    bundles = zip(range(len(keyWordsStrList)),keyWordsStrList)
    bundles = list(bundles)
    bundles = np.array(bundles)
    np.apply_along_axis(func1d=line2count,axis=1,arr=bundles)

    unNormalizedMatrix = np.matmul(keyWordsStrMatrix, matrix)
    keyWordsStrMatrix = unNormalizedMatrix / normalizer

    return keyWordsStrMatrix

def sentence2vector(sentence,word2id,matrix):
    keyWordsStr = dictSentence2KeyWordsStr[sentence]
    sentenceVectorMatrix = get_keyWordsSentence_matrix(word2id,matrix,keyWordsStrList=[keyWordsStr])
    if isCodingAndDebug and False:
        print(sentenceVectorMatrix)
    return sentenceVectorMatrix

def getTopKids(matrix_q,matrix_v,TopK=5):
    import util
    cosMatrix = util.MatrixCos(matrix_q,matrix_v,isTrimZeros=False)
    sortedCos = np.sort(cosMatrix)[:,::-1][:,:TopK]
    sortedIDS = np.argsort(cosMatrix)[:,::-1][:,:TopK]
    if isCodingAndDebug:
        print("余弦相似度：",sortedCos)
        print("对应的公式id",sortedIDS)
    return sortedIDS

def singleSentenceSearch(sentence,word2id,matrix,keyWordsStrMatrix,TopK=5):
    sentencesMatrix = sentence2vector(sentence,word2id,matrix)
    sortedIDS = getTopKids(sentencesMatrix, keyWordsStrMatrix,TopK=TopK)
    candidates = np.vectorize(lambda idx:dictID2Sentence[idx])(sortedIDS)
    if isCodingAndDebug :
        print("input sentence: ",sentence)
        print("candidates TopK=",TopK)
        for i in range(candidates.shape[-1]):
            print(candidates[0][i])
    return candidates

if __name__ == "__main__":
    # 获取词向量，词表
    words,word2id,matrix = get_words_matrix()
    keyWordsStrMatrix = get_keyWordsSentence_matrix(word2id, matrix)
    topK = 5
    #x=\frac{\negative*b\pm\sqrt{b^{2}-4*a*c}}{2*a}
    sentences = list(dictSentence2KeyWordsStr.keys())[:2]
    for sentence in sentences:
        time_s = time.time()
        singleSentenceSearch(sentence,word2id,matrix,keyWordsStrMatrix,TopK=topK)
        time_e = time.time()
        print("耗时：{}毫秒".format((time_e-time_s)*1000))

    active = True
    prompt = '\n请输入TopK数字，表示选择多少个最相似的公式,输入quit将会退出。:'
    while prompt:
        try:
            message = input(prompt)
            if message == 'quit':
                active = False
                prompt = False
            else:
                topK = int(message)
                prompt = False
        except:pass

    prompt = "\n请输入待查询的公式，输入'quit'将会退出。:"
    while active:
        message = input(prompt)
        if message == 'quit':
            active = False
        else:
            sentence = message
            try:
                time_s = time.time()
                singleSentenceSearch(sentence, word2id, matrix, keyWordsStrMatrix, TopK=topK)
                time_e = time.time()
                print("耗时：{}毫秒".format((time_e - time_s) * 1000))
            except:print('输入公式不在，已知切分的公式集合中。')
